"""
电商产品搜索示例 - 使用SelfQueryRetriever

这个示例展示了如何使用SelfQueryRetriever来实现带过滤条件的电商产品搜索。
用户可以输入自然语言查询，系统会自动解析查询中的过滤条件（如价格范围、品牌、评分等）
并结合语义搜索来找到最匹配的产品。
"""
from langchain.chains.query_constructor.schema import AttributeInfo
from langchain.retrievers import SelfQueryRetriever
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from models import get_ds_model_client, get_ollama_embeddings_client

# 电商产品检索示例
# 定义产品元数据
product_metadata = [
    AttributeInfo(
        name="price",
        description="产品价格，单位为元",
        type="float",
    ),
    AttributeInfo(
        name="brand",
        description="产品品牌",
        type="string",
    ),
    AttributeInfo(
        name="rating",
        description="用户评分，1-5星",
        type="integer",
    ),
    AttributeInfo(
        name="category",
        description="产品类别",
        type="string",
    ),
    AttributeInfo(
        name="stock",
        description="库存数量",
        type="integer",
    )
]

# 创建示例产品数据
example_products = [
    Document(
        page_content="Apple MacBook Pro 14英寸笔记本电脑，配备M3 Pro芯片，16GB内存，512GB SSD",
        metadata={"price": 12999.0, "brand": "Apple", "rating": 4.8, "category": "笔记本电脑", "stock": 15}
    ),
    Document(
        page_content="Apple iPhone 15 Pro Max 256GB 钛金属原色",
        metadata={"price": 9999.0, "brand": "Apple", "rating": 4.9, "category": "手机", "stock": 30}
    ),
    Document(
        page_content="Samsung Galaxy S24 Ultra 256GB 幻影黑",
        metadata={"price": 8999.0, "brand": "Samsung", "rating": 4.7, "category": "手机", "stock": 20}
    ),
    Document(
        page_content="Sony WH-1000XM5 无线降噪耳机",
        metadata={"price": 2999.0, "brand": "Sony", "rating": 4.6, "category": "耳机", "stock": 10}
    ),
    Document(
        page_content="Bose QuietComfort Ultra 无线降噪耳机",
        metadata={"price": 3299.0, "brand": "Bose", "rating": 4.7, "category": "耳机", "stock": 5}
    ),
    Document(
        page_content="Dell XPS 15 笔记本电脑，12代i7处理器，16GB内存，1TB SSD",
        metadata={"price": 9999.0, "brand": "Dell", "rating": 4.5, "category": "笔记本电脑", "stock": 8}
    ),
    Document(
        page_content="华为 MateBook X Pro 2024款 14.2英寸笔记本电脑",
        metadata={"price": 7999.0, "brand": "华为", "rating": 4.6, "category": "笔记本电脑", "stock": 12}
    ),
    Document(
        page_content="小米 14 Pro 12GB+256GB 黑色",
        metadata={"price": 4999.0, "brand": "小米", "rating": 4.5, "category": "手机", "stock": 25}
    )
]

def create_ecommerce_retriever():
    """创建电商产品检索器"""
    # 初始化LLM
    try:
        llm = get_ds_model_client()
    except Exception as e:
        print(f"无法连接到DeepSeek模型，将使用本地Ollama模型: {e}")
        from models import get_ds_local_model_client
        llm = get_ds_local_model_client()

    # 创建向量存储
    vectorstore = Chroma.from_documents(
        documents=example_products,
        embedding=get_ollama_embeddings_client()
    )

    # 创建SelfQueryRetriever
    retriever = SelfQueryRetriever.from_llm(
        llm=llm,
        vectorstore=vectorstore,
        document_contents="产品描述",
        metadata_field_info=product_metadata,
        verbose=True
    )

    return retriever


def search_products(query, retriever):
    """搜索产品"""
    print(f"\n搜索查询: {query}")
    documents = retriever.get_relevant_documents(query)
    print(f"找到 {len(documents)} 个匹配产品:")
    for i, doc in enumerate(documents, 1):
        print(f"\n产品 {i}:")
        print(f"描述: {doc.page_content}")
        print(f"价格: {doc.metadata['price']}元")
        print(f"品牌: {doc.metadata['brand']}")
        print(f"评分: {doc.metadata['rating']}星")
        print(f"类别: {doc.metadata['category']}")
        print(f"库存: {doc.metadata['stock']}件")
    return documents


if __name__ == "__main__":
    print("===== 电商产品搜索系统 =====")
    retriever = create_ecommerce_retriever()

    # 预设查询示例
    print("\n===== 预设查询示例 =====")
    preset_queries = [
        "价格低于10000元的Apple笔记本电脑",
        "评分4.7星以上的手机",
        "价格在2000-3500元之间的无线耳机",
        "库存大于10的华为产品"
    ]
    
    for query in preset_queries:
        search_products(query, retriever)
        input("\n按Enter键继续...")
